import os
import math
import jieba
import collections

import torch
import torch.nn as nn


def gpu_setup(use_gpu, gpu_id):
    """GPU设置"""
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    if torch.cuda.is_available() and use_gpu:
        print('cuda available with GPU:', torch.cuda.get_device_name(0))
        device = torch.device("cuda")
    else:
        print('cuda not available')
        device = torch.device("cpu")
    return device


def xavier_init_weights(m):
    """用Xavier的方法初始化权重"""
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
    if type(m) == nn.GRU:
        for param in m._flat_weights_names:
            if "weight" in param:
                nn.init.xavier_uniform_(m._parameters[param])


def grad_clipping(model, max_norm):
    """裁剪梯度"""
    if isinstance(model, nn.Module):
        params = [p for p in model.parameters() if p.requires_grad]
    else:
        params = model.params
    nn.utils.clip_grad_norm_(params, max_norm=max_norm, norm_type=2)


def split_ch(text, method='word'):
    # 先去掉空格后，再分词
    text = text.replace(' ', '')
    if method == 'char':
        out = [char for i, char in enumerate(text)]
    elif method == 'word':
        seg_list = jieba.cut(text)
        zh_words = " ".join(seg_list)
        out = zh_words.split()
    else:
        print("没有这种分词方法！")
        raise
    return out


def bleu(pred_seq, label_seq, k, method):
    """计算BLEU值"""
    pred_tokens = pred_seq.split()
    label_tokens = split_ch(label_seq, method)
    len_pred, len_label = len(pred_tokens), len(label_tokens)
    score = math.exp(min(0, 1 - len_label / len_pred))
    for n in range(1, k + 1):
        num_matches, label_subs = 0, collections.defaultdict(int)
        for i in range(len_label - n + 1):
            label_subs[' '.join(label_tokens[i: i + n])] += 1
        for i in range(len_pred - n + 1):
            if label_subs[' '.join(pred_tokens[i: i + n])] > 0:
                num_matches += 1
                label_subs[' '.join(pred_tokens[i: i + n])] -= 1
        if (len_pred - n + 1) == 0:  # 防止报错
            score *= math.pow(num_matches / 1e-6, math.pow(0.5, n))
        else:
            score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))
    return score
